"""
特征工程模块
负责特征提取、创建和转换
"""

import pandas as pd
import numpy as np
import warnings
from sklearn.preprocessing import LabelEncoder, StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer

# 忽略警告
warnings.filterwarnings('ignore')


def extract_title(df):
    """
    从姓名中提取头衔

    参数:
        df (pd.DataFrame): 原始数据框

    返回:
        pd.DataFrame: 添加了头衔特征的数据框
    """
    print("=" * 50)
    print("提取头衔特征")
    print("=" * 50)
    
    df_processed = df.copy()

    # 检查Name列是否存在
    if 'Name' not in df_processed.columns:
        print("警告: 数据框中没有Name列，无法提取头衔")
        return df_processed

    # 从姓名中提取头衔
    df_processed['Title'] = df_processed['Name'].str.extract(' ([A-Za-z]+)\.', expand=False)

    # 统计头衔出现频率
    title_counts = df_processed['Title'].value_counts()
    print(f"提取到的头衔及数量: {title_counts.to_dict()}")

    # 合并稀有头衔
    rare_titles = title_counts[title_counts < 10].index
    df_processed['Title'] = df_processed['Title'].replace(rare_titles, 'Rare')

    # 创建头衔映射字典，便于后续分析
    title_mapping = {
        'Mr': 1, 'Miss': 2, 'Mrs': 3, 'Master': 4, 'Rare': 5
    }
    df_processed['Title_Code'] = df_processed['Title'].map(title_mapping).fillna(0).astype(int)

    print("头衔提取完成")
    print(f"合并后的头衔分布: {df_processed['Title'].value_counts().to_dict()}")

    return df_processed


def create_family_features(df):
    """
    创建家庭相关特征

    参数:
        df (pd.DataFrame): 原始数据框

    返回:
        pd.DataFrame: 添加了家庭特征的数据框
    """
    print("=" * 50)
    print("创建家庭特征")
    print("=" * 50)
    
    df_processed = df.copy()

    # 检查必要的列是否存在
    required_cols = ['SibSp', 'Parch']
    missing_cols = [col for col in required_cols if col not in df_processed.columns]
    if missing_cols:
        print(f"警告: 数据框中缺少以下列: {missing_cols}，无法创建家庭特征")
        return df_processed

    # 家庭规模 = 兄弟姐妹数 + 父母子女数 + 1（本人）
    df_processed['FamilySize'] = df_processed['SibSp'] + df_processed['Parch'] + 1

    # 是否独自旅行（家庭规模 = 1 则为 1，否则为 0）
    df_processed['IsAlone'] = (df_processed['FamilySize'] == 1).astype(int)

    # 家庭规模分组
    df_processed['FamilySizeGroup'] = pd.cut(
        df_processed['FamilySize'], 
        bins=[0, 1, 3, 5, 100], 
        labels=['alone', 'small', 'medium', 'large']
    )
    
    # 创建家庭类型特征
    df_processed['FamilyType'] = 'single'
    df_processed.loc[(df_processed['SibSp'] > 0) & (df_processed['Parch'] == 0), 'FamilyType'] = 'sibling_only'
    df_processed.loc[(df_processed['SibSp'] == 0) & (df_processed['Parch'] > 0), 'FamilyType'] = 'parent_child_only'
    df_processed.loc[(df_processed['SibSp'] > 0) & (df_processed['Parch'] > 0), 'FamilyType'] = 'both'

    print("家庭特征创建完成")
    print(f"家庭规模分布: {df_processed['FamilySize'].value_counts().sort_index().to_dict()}")
    print(f"家庭类型分布: {df_processed['FamilyType'].value_counts().to_dict()}")

    return df_processed


def extract_ticket_features(df):
    """
    从船票号中提取特征

    参数:
        df (pd.DataFrame): 原始数据框

    返回:
        pd.DataFrame: 添加了船票特征的数据框
    """
    print("=" * 50)
    print("提取船票特征")
    print("=" * 50)
    
    df_processed = df.copy()

    # 检查Ticket列是否存在
    if 'Ticket' not in df_processed.columns:
        print("警告: 数据框中没有Ticket列，无法提取船票特征")
        return df_processed

    # 提取船票前缀
    df_processed['Ticket_Prefix'] = df_processed['Ticket'].apply(lambda x: str(x).split()[0] if len(str(x).split()) > 1 else 'None')
    
    # 统计船票前缀频率
    prefix_counts = df_processed['Ticket_Prefix'].value_counts()
    
    # 合并稀有前缀
    rare_prefixes = prefix_counts[prefix_counts < 5].index
    df_processed['Ticket_Prefix'] = df_processed['Ticket_Prefix'].replace(rare_prefixes, 'Rare')
    
    # 提取船票号码
    df_processed['Ticket_Number'] = df_processed['Ticket'].apply(lambda x: str(x).split()[-1] if str(x).split()[-1].isdigit() else 0)
    df_processed['Ticket_Number'] = pd.to_numeric(df_processed['Ticket_Number'], errors='coerce').fillna(0)
    
    # 船票长度特征
    df_processed['Ticket_Length'] = df_processed['Ticket'].apply(lambda x: len(str(x)))

    print("船票特征提取完成")
    print(f"船票前缀分布: {df_processed['Ticket_Prefix'].value_counts().head(10).to_dict()}")

    return df_processed


def bin_numeric_features(df):
    """
    对数值特征进行分箱处理

    参数:
        df (pd.DataFrame): 原始数据框

    返回:
        pd.DataFrame: 添加了分箱特征的数据框
    """
    print("=" * 50)
    print("数值特征分箱")
    print("=" * 50)
    
    df_processed = df.copy()

    # 年龄分箱
    if 'Age' in df_processed.columns:
        df_processed['AgeBin'] = pd.cut(
            df_processed['Age'], 
            bins=[0, 12, 18, 35, 60, 100], 
            labels=['child', 'teenager', 'young_adult', 'adult', 'senior']
        )
        print(f"年龄分箱分布: {df_processed['AgeBin'].value_counts().to_dict()}")

    # 票价分箱
    if 'Fare' in df_processed.columns:
        # 处理可能的极端值
        fare_q1 = df_processed['Fare'].quantile(0.25)
        fare_q3 = df_processed['Fare'].quantile(0.75)
        fare_iqr = fare_q3 - fare_q1
        upper_bound = fare_q3 + 1.5 * fare_iqr
        
        # 将极端值替换为上限值
        df_processed['Fare'] = np.where(df_processed['Fare'] > upper_bound, upper_bound, df_processed['Fare'])
        
        # 分箱
        df_processed['FareBin'] = pd.qcut(
            df_processed['Fare'], 
            q=5, 
            labels=['very_low', 'low', 'medium', 'high', 'very_high'],
            duplicates='drop'
        )
        print(f"票价分箱分布: {df_processed['FareBin'].value_counts().to_dict()}")

    print("数值特征分箱完成")

    return df_processed


def create_interaction_features(df):
    """
    创建交互特征

    参数:
        df (pd.DataFrame): 原始数据框

    返回:
        pd.DataFrame: 添加了交互特征的数据框
    """
    print("=" * 50)
    print("创建交互特征")
    print("=" * 50)
    
    df_processed = df.copy()

    # 性别和船舱等级的交互特征
    if 'Sex' in df_processed.columns and 'Pclass' in df_processed.columns:
        df_processed['Sex_Pclass'] = df_processed['Sex'].astype(str) + "_" + df_processed['Pclass'].astype(str)
        print(f"性别和船舱等级交互特征创建完成，唯一值数量: {df_processed['Sex_Pclass'].nunique()}")

    # 年龄和船舱等级的交互特征
    if 'AgeBin' in df_processed.columns and 'Pclass' in df_processed.columns:
        df_processed['Age_Pclass'] = df_processed['AgeBin'].astype(str) + "_" + df_processed['Pclass'].astype(str)
        print(f"年龄和船舱等级交互特征创建完成，唯一值数量: {df_processed['Age_Pclass'].nunique()}")

    # 家庭规模和船舱等级的交互特征
    if 'FamilySizeGroup' in df_processed.columns and 'Pclass' in df_processed.columns:
        df_processed['FamilySize_Pclass'] = df_processed['FamilySizeGroup'].astype(str) + "_" + df_processed['Pclass'].astype(str)
        print(f"家庭规模和船舱等级交互特征创建完成，唯一值数量: {df_processed['FamilySize_Pclass'].nunique()}")

    print("交互特征创建完成")
    return df_processed


def encode_categorical_features(df, encoding_method='label'):
    """
    对类别特征进行编码

    参数:
        df (pd.DataFrame): 原始数据框
        encoding_method (str): 编码方法，'label'或'onehot'

    返回:
        tuple: (处理后的数据框, 编码器字典)
    """
    print("=" * 50)
    print("类别特征编码")
    print("=" * 50)
    
    df_processed = df.copy()
    encoders = {}

    # 需要编码的类别特征
    categorical_features = ['Sex', 'Embarked', 'Title', 'Cabin_Letter', 'AgeBin', 'FareBin', 'FamilySizeGroup', 
                           'FamilyType', 'Ticket_Prefix', 'Sex_Pclass', 'Age_Pclass', 'FamilySize_Pclass']

    # 只处理存在于数据框中的特征
    features_to_encode = [feature for feature in categorical_features if feature in df_processed.columns]
    
    if not features_to_encode:
        print("没有需要编码的类别特征")
        return df_processed, encoders

    if encoding_method == 'label':
        # 标签编码
        for feature in features_to_encode:
            le = LabelEncoder()
            df_processed[feature] = le.fit_transform(df_processed[feature])
            encoders[feature] = le
            print(f"已对特征 '{feature}' 进行标签编码，类别数量: {len(le.classes_)}")
    
    elif encoding_method == 'onehot':
        # 独热编码
        # 使用ColumnTransformer进行独热编码
        ct = ColumnTransformer(
            transformers=[('encoder', OneHotEncoder(drop='first', sparse=False), features_to_encode)],
            remainder='passthrough'
        )
        
        # 获取编码后的数据
        encoded_data = ct.fit_transform(df_processed)
        
        # 获取新的列名
        onehot_features = []
        for i, feature in enumerate(features_to_encode):
            if i < len(ct.transformers_[0][1].categories_):
                for category in ct.transformers_[0][1].categories_[i][1:]:  # 跳过第一个类别
                    onehot_features.append(f"{feature}_{category}")
        
        # 获取未编码的列名
        other_features = [col for col in df_processed.columns if col not in features_to_encode]
        
        # 创建新的DataFrame
        df_processed = pd.DataFrame(encoded_data, columns=onehot_features + other_features, index=df_processed.index)
        encoders['column_transformer'] = ct
        
        print(f"已对 {len(features_to_encode)} 个特征进行独热编码，生成 {len(onehot_features)} 个新特征")
    
    else:
        print(f"警告: 不支持的编码方法 '{encoding_method}'，将使用标签编码")
        for feature in features_to_encode:
            le = LabelEncoder()
            df_processed[feature] = le.fit_transform(df_processed[feature])
            encoders[feature] = le
            print(f"已对特征 '{feature}' 进行标签编码，类别数量: {len(le.classes_)}")

    print("类别特征编码完成")
    return df_processed, encoders


def select_features(df, exclude_cols=None, include_cols=None):
    """
    特征选择，删除无关特征或选择特定特征

    参数:
        df (pd.DataFrame): 原始数据框
        exclude_cols (list): 要排除的列名列表
        include_cols (list): 要包含的列名列表（如果提供，则只保留这些列）

    返回:
        pd.DataFrame: 选择特征后的数据框
    """
    print("=" * 50)
    print("特征选择")
    print("=" * 50)
    
    if exclude_cols is None:
        exclude_cols = ['PassengerId', 'Name', 'Ticket', 'Cabin']

    df_processed = df.copy()
    original_cols = list(df_processed.columns)
    
    # 如果指定了包含列，则只保留这些列
    if include_cols is not None:
        # 确保包含的列存在于数据框中
        valid_include_cols = [col for col in include_cols if col in df_processed.columns]
        if len(valid_include_cols) < len(include_cols):
            missing_cols = set(include_cols) - set(valid_include_cols)
            print(f"警告: 以下指定的列不存在于数据框中: {missing_cols}")
        
        df_processed = df_processed[valid_include_cols]
        print(f"特征选择完成，保留了 {len(valid_include_cols)} 个指定特征")
        return df_processed
    
    # 否则，排除指定的列
    # 删除无关特征
    for col in exclude_cols:
        if col in df_processed.columns:
            df_processed = df_processed.drop(col, axis=1)
            print(f"已排除特征: {col}")
    
    remaining_cols = list(df_processed.columns)
    print(f"特征选择完成，原始特征数: {len(original_cols)}，剩余特征数: {len(remaining_cols)}")
    print(f"剩余特征: {remaining_cols}")

    return df_processed


def scale_features(X_train, X_val=None, method='standard'):
    """
    特征标准化

    参数:
        X_train (pd.DataFrame): 训练集特征
        X_val (pd.DataFrame, optional): 验证集特征
        method (str): 标准化方法，'standard'或'minmax'

    返回:
        tuple: (标准化后的训练集, 标准化后的验证集, 标准化器)
    """
    print("=" * 50)
    print("特征标准化")
    print("=" * 50)
    
    if method == 'standard':
        scaler = StandardScaler()
        print("使用标准化(StandardScaler)方法")
    elif method == 'minmax':
        from sklearn.preprocessing import MinMaxScaler
        scaler = MinMaxScaler()
        print("使用最小-最大标准化(MinMaxScaler)方法")
    else:
        print(f"警告: 不支持的标准化方法 '{method}'，将使用标准化(StandardScaler)方法")
        scaler = StandardScaler()

    # 对训练集进行标准化
    X_train_scaled = scaler.fit_transform(X_train)
    X_train_scaled = pd.DataFrame(X_train_scaled, columns=X_train.columns, index=X_train.index)

    # 如果有验证集，也进行标准化
    X_val_scaled = None
    if X_val is not None:
        X_val_scaled = scaler.transform(X_val)
        X_val_scaled = pd.DataFrame(X_val_scaled, columns=X_val.columns, index=X_val.index)
        print("验证集已标准化")

    print("特征标准化完成")
    return X_train_scaled, X_val_scaled, scaler


def apply_feature_engineering_pipeline(df, is_train=True, encoders=None, scaler=None, encoding_method='label', scaling_method='standard'):
    """
    应用完整的特征工程流程

    参数:
        df (pd.DataFrame): 原始数据框
        is_train (bool): 是否为训练数据
        encoders (dict, optional): 训练好的编码器
        scaler (StandardScaler, optional): 训练好的标准化器
        encoding_method (str): 编码方法，'label'或'onehot'
        scaling_method (str): 标准化方法，'standard'或'minmax'

    返回:
        tuple: (处理后的数据框, 编码器, 标准化器)
    """
    print("=" * 50)
    print("应用特征工程流程")
    print("=" * 50)
    
    df_processed = df.copy()
    
    try:
        # 1. 提取头衔
        df_processed = extract_title(df_processed)

        # 2. 创建家庭特征
        df_processed = create_family_features(df_processed)
        
        # 3. 提取船票特征
        df_processed = extract_ticket_features(df_processed)

        # 4. 数值特征分箱
        df_processed = bin_numeric_features(df_processed)
        
        # 5. 创建交互特征
        df_processed = create_interaction_features(df_processed)

        # 6. 类别特征编码
        if is_train:
            df_processed, encoders = encode_categorical_features(df_processed, encoding_method)
        else:
            # 使用训练好的编码器
            if encoders is None:
                print("警告: 测试数据编码时未提供编码器，将跳过编码步骤")
            else:
                if encoding_method == 'label':
                    for feature, encoder in encoders.items():
                        if feature in df_processed.columns:
                            # 处理测试集中可能存在的新类别
                            unique_values = set(df_processed[feature].unique())
                            encoder_classes = set(encoder.classes_)
                            new_values = unique_values - encoder_classes

                            if new_values:
                                # 将新类别映射为最常见的类别
                                for val in new_values:
                                    df_processed.loc[df_processed[feature] == val, feature] = encoder.classes_[0]
                                print(f"特征 '{feature}' 中发现新类别 {new_values}，已映射为最常见的类别")

                            df_processed[feature] = encoder.transform(df_processed[feature])
                            print(f"已对特征 '{feature}' 应用编码")
                elif encoding_method == 'onehot' and 'column_transformer' in encoders:
                    ct = encoders['column_transformer']
                    # 获取编码后的数据
                    encoded_data = ct.transform(df_processed)
                    
                    # 获取新的列名
                    categorical_features = ct.transformers_[0][2]
                    onehot_features = []
                    for i, feature in enumerate(categorical_features):
                        if i < len(ct.transformers_[0][1].categories_):
                            for category in ct.transformers_[0][1].categories_[i][1:]:  # 跳过第一个类别
                                onehot_features.append(f"{feature}_{category}")
                    
                    # 获取未编码的列名
                    other_features = [col for col in df_processed.columns if col not in categorical_features]
                    
                    # 创建新的DataFrame
                    df_processed = pd.DataFrame(encoded_data, columns=onehot_features + other_features, index=df_processed.index)
                    print(f"已应用独热编码，生成 {len(onehot_features)} 个新特征")

        # 7. 特征选择
        df_processed = select_features(df_processed)

        # 8. 特征标准化（如果不是目标变量）
        if 'Survived' in df_processed.columns:
            X = df_processed.drop('Survived', axis=1)
            y = df_processed['Survived']

            if is_train:
                X_scaled, _, scaler = scale_features(X, method=scaling_method)
                df_processed = pd.concat([X_scaled, y], axis=1)
            else:
                if scaler is None:
                    print("警告: 测试数据标准化时未提供标准化器，将跳过标准化步骤")
                    df_processed = pd.concat([X, y], axis=1) if 'Survived' in df.columns else X
                else:
                    X_scaled, _, _ = scale_features(X, method=scaling_method)
                    df_processed = pd.concat([X_scaled, y], axis=1) if 'Survived' in df.columns else X_scaled
        else:
            if is_train:
                df_processed, _, scaler = scale_features(df_processed, method=scaling_method)
            else:
                if scaler is None:
                    print("警告: 测试数据标准化时未提供标准化器，将跳过标准化步骤")
                else:
                    df_processed, _, _ = scale_features(df_processed, method=scaling_method)

        print("特征工程流程应用完成")
        return df_processed, encoders, scaler
    
    except Exception as e:
        print(f"特征工程流程应用出错: {e}")
        raise
